from langChain.output_parsers import YamlOutputParser, XMLOutputParser
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from pydantic import BaseModel, Field

from langChain.config import model

# 一 、 JSON格式
# 定义预期的数据结构
class joke(BaseModel):
    setup: str = Field(description="设置笑话的问题")
    punchline: str = Field(description=" 解决笑话的答案")


# JsonOutputParser 是 LangChain 中用于处理模型输出的一个解析器，它主要用于将语言模型的文本响应转换为符合特定JSON格式的结构化数据。
# 按照指定格式输出json,指定格式就是 joke
parser = JsonOutputParser(pydantic_object=joke)

# # get_format_instructions ：返回一系列格式化文本指令，这些指令用于指导语言模型如何格式化其响应以符合预设的JSON模式。
# print('get_format_instructions 格式化文本指令：', parser.get_format_instructions(), '\n--------------------------------')

# # 提示词模板提前初始化变量
# # 1。以前提示词变量是 动态赋值，提示词模板中的变量都是在运行时动态生成的。
# prompt_1 = PromptTemplate(
#     template="请按用户指令生成一个笑话。\n{format_instructions}\n{query}\n",
#     input_variables=["query", "format_instructions"],
# )
# final_template_1 = prompt_1.format(format_instructions=parser.get_format_instructions(), query="告诉我一个笑话。")
# print(f"提示词变量动态赋值，最终生成的提示词 ：\n {final_template_1} \n--------------------------------")

# # 2。现在提示词变量是 静态赋值，提示词模板中的变量可以在模板实例化时确定。
# # partial_variables :字典类型, 指定那些在模板实例化时就需要确定的变量
# prompt_2 = PromptTemplate(
#     template="请按用户指令生成一个笑话。\n{format_instructions}\n{query}\n",
#     input_variables=["query", "format_instructions"],
#     partial_variables={"format_instructions": parser.get_format_instructions()},
# )
#
# # 这里不用指定format_instructions值，因为它在提示词模板生成时就已经指定了
# final_template_2 = prompt_2.format(query="告诉我一个笑话。")
# print(f"提示词变量静态赋值，最终生成的提示词 ：\n {final_template_2}  \n--------------------------------")

# chain = prompt_2 | model | parser
# # result = chain.invoke({"query": " Tell me a joke.","format_instructions": parser.get_format_instructions()})
# result = chain.invoke({"query": " Tell me a joke."})
# print('AI生成的笑话 ：', result, '\n--------------------------------')

# # 流式处理
# stream_result = chain.stream({"query": " Tell me a joke."})
# for chunk in stream_result:
#     print(chunk)

# # 二 、 xml格式。需要安装 defusedxml
# parser = XMLOutputParser()
# prompt_3 = PromptTemplate(
#     template="请按用户指令完成任务。\n{format_instructions}\n{query}\n",
#     input_variables=["query", "format_instructions"],
#     partial_variables={"format_instructions": parser.get_format_instructions()},
# )
# chain = prompt_3 | model
# result = chain.invoke({"query": " 生成周星驰的简化电影列表，按照时间降序"})
# print('周星驰的电影列表 ：\n', parser.parse(result.content), '\n--------------------------------')

# # 三 、 yaml格式。
# # 定义预期的数据结构
# class joke(BaseModel):
#     setup: str = Field(description="设置笑话的问题")
#     punchline: str = Field(description=" 解决笑话的答案")
#
#
# # JsonOutputParser 是 LangChain 中用于处理模型输出的一个解析器，它主要用于将语言模型的文本响应转换为符合特定JSON格式的结构化数据。
# # 按照指定格式输出json,指定格式就是 joke
# parser = YamlOutputParser(pydantic_object=joke)
#
# prompt_6 = PromptTemplate(
#     template="请按用户指令生成一个笑话。\n{format_instructions}\n{query}\n",
#     input_variables=["query", "format_instructions"],
#     partial_variables={"format_instructions": parser.get_format_instructions()},
# )
# chain = prompt_6 | model
#
# result = chain.invoke({"query": " Tell me a joke."})
# print('AI生成的笑话 ：', parser.parse(result.content), '\n--------------------------------')
